from jqdatasdk import *
# from jqdata import *
from jq_utils import *
import datetime
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
pd.reset_option('display.float_format')

sdk_auth()

def pdate(dstr):
    return datetime.datetime.strptime(dstr, '%Y-%m-%d').date()

def get_index_codes(codes = "000300.XSHG"):
    ret = []
    if type(codes) == list:
        for code in codes:
            ret += get_index_stocks(code, date=None)
    else:
        ret = get_index_stocks(codes)
    return ret

def get_hold_df(trade_days, link_ids = [310001, 310002]):
    df_map = {}
    for d in trade_days:
        print(d)
        pdt = d if type(d) == datetime.date else pdate(d)
        df_d=finance.run_query(query(finance.STK_HK_HOLD_INFO).filter(
                    (finance.STK_HK_HOLD_INFO.link_id.in_(link_ids)) &
                   (finance.STK_HK_HOLD_INFO.day == pdt)
        ))
        print(df_d)
        df_map[d] = df_d
    return df_map

def get_inc_hold(hdf, sd, ed, codes = None):
    sdf = hdf[sd]
    edf = hdf[ed]
    res = []
    y1st = pdate('2021-01-04')
    for i, row in edf.iterrows():
        code = row['code']
        if codes and code not in codes:
            continue
        py1 = get_price(code, start_date=y1st, end_date=y1st, frequency='daily', fields=[ 'close'])['close'].values[0]
        pst = get_price(code, start_date=sd, end_date=sd, frequency='daily', fields=[ 'close'])['close'].values[0]
        ped = get_price(code, start_date=ed, end_date=ed, frequency='daily', fields=[ 'close'])['close'].values[0]
        srs = sdf.loc[sdf.code == code, 'share_ratio'].values[0]
        sre = edf.loc[edf.code == code, 'share_ratio'].values[0]
        p_lift = (ped / pst) - 1
        py_lift = (ped / py1) - 1
        lift = 1 if srs == 0 else (sre / srs) - 1
        if sre > 0 and sre > srs :#and lift > 0.1 and p_lift < 0.1:
            res.append((row['name'],row['code'] , srs, sre, sre-srs, lift, pst, ped, p_lift,py1, py_lift))
    return pd.DataFrame(res, columns = ['name', 'code', 'start_sr', 'end_sr', 'sr_diff', 'sr_lift', 'pst', 'ped', 'p_lift', 'py1', 'py_lift'])


zz800 = get_index_codes([ "000905.XSHG", "000300.XSHG"])
latest_trade_days = get_trade_days(count = 10)
last_hold_df = get_hold_df(latest_trade_days)

inc_df = get_inc_hold(last_hold_df, latest_trade_days[-5], latest_trade_days[-2], zz800)
# inc_df[inc_df.sr_diff > 0].sort_values(['sr_diff'], ascending = False)
print(inc_df)
